#############################################################################################
#############################################################################################
#
#   Open Source License/Disclaimer, Forecast Systems Laboratory NOAA/OAR/GSD,
#   325 Broadway Boulder, CO 80305
#
#   This software is distributed under the Open Source Definition, which may be
#   found at http://www.opensource.org/.
#   In particular, redistribution and use in source and binary forms, with or
#   without modification, are permitted provided that the following conditions are met:
#
#   - Redistributions of source code must retain this notice, this list of
#     conditions and the following disclaimer.
#
#   - Redistributions in binary form must provide access to this notice, this
#     list of conditions and the following disclaimer, and the underlying source code.
#
#   - All modifications to this software must be clearly documented, and are
#     solely the responsibility of the agent making the modifications.
#
#   - If significant modifications or enhancements are made to this software,
#     the GSD Software Policy Manager (softwaremgr.fsl@noaa.gov) should be notified.
#
#   THIS SOFTWARE AND ITS DOCUMENTATION ARE IN THE PUBLIC DOMAIN AND ARE
#   FURNISHED "AS IS." THE AUTHORS, THE UNITED STATES GOVERNMENT, ITS INSTRUMENTALITIES,
#   OFFICERS, EMPLOYEES, AND AGENTS MAKE NO WARRANTY, EXPRESS OR IMPLIED, AS TO
#   THE USEFULNESS OF THE SOFTWARE AND DOCUMENTATION FOR ANY PURPOSE. THEY ASSUME
#   NO RESPONSIBILITY (1) FOR THE USE OF THE SOFTWARE AND DOCUMENTATION; OR (2) TO PROVIDE
#   TECHNICAL SUPPORT TO USERS.
#############################################################################################
#############################################################################################

import numpy as np

from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.colors import cnames
from matplotlib import animation


def movie(x_t, N_trajectories=1):
    
    # Set up figure & 3D axis for animation
    fig = plt.figure()
    ax = fig.add_axes([0, 0, 1, 1], projection='3d')
    ax.axis('off')
    
    # choose a different color for each trajectory
    colors = plt.cm.jet(np.linspace(0, 1, N_trajectories))
    
    # set up lines and points
    lines = sum([ax.plot([], [], [], '-', c=c)
                for c in colors], [])
    pts = sum([ax.plot([], [], [], 'o', c=c)
            for c in colors], [])
    
    # prepare the axes limits
    ax.set_xlim((-25, 25))
    ax.set_ylim((-35, 35))
    ax.set_zlim((5, 55))
    
    # set point-of-view: specified by (altitude degrees, azimuth degrees)
    ax.view_init(30, 0)
    
    # initialization function: plot the background of each frame
    def init():
        for line, pt in zip(lines, pts):
            line.set_data([], [])
            line.set_3d_properties([])
    
            pt.set_data([], [])
            pt.set_3d_properties([])
        return lines + pts
    
    # animation function.  This will be called sequentially with the frame number
    def animate(i):
        # we'll step two time-steps per frame.  This leads to nice results.
        i = (2 * i) % x_t.shape[1]
    
        for line, pt, xi in zip(lines, pts, x_t):
            x, y, z = xi[:i].T
            line.set_data(x, y)
            line.set_3d_properties(z)
            
    
            pt.set_data(x[-1:], y[-1:])
            pt.set_3d_properties(z[-1:])
    
        ax.view_init(30, 0.3 * i)
        fig.canvas.draw()
        return lines + pts
    
    # instantiate the animator.
    anim = animation.FuncAnimation(fig, animate, init_func=init,
                                frames=500, interval=30, blit=True)
    
    # Save as mp4. This requires mplayer or ffmpeg to be installed
    #anim.save('lorentz_attractor.mp4', fps=15, extra_args=['-vcodec', 'libx264'])
    
    plt.show()